#!/usr/bin/perl

# This script is the workhorse of the HippDB pipeline. Given a PDB code and optionally 
# two chains, it figures out what the next step is for that input and calls a job to 
# perform it. 

use strict;
use warnings;

use DBI;

my $usage = "$0 PDB_code [chain_1 chain_2]\n";

my $pdb = shift;
my $chain1 = shift;
my $chain2 = shift;

defined $pdb
	or die "ERROR: No PDB code passed.\n$usage";

# Load the module full of tests, like "was the PDB file downloaded?"
use lib './scripts';
use check_status;

# Load module to submit PBS jobs
# http://search.cpan.org/~kwmak/PBS-Client/lib/PBS/Client.pm
use lib '/scratch/cmb433/pdb_interact/dependencies/PBS-Client-0.09/lib';
use PBS::Client;

my $pbs = PBS::Client->new;

# E-mail address for sending job reports
my @emails = ['cmb433@nyu.edu'];

######
### Are we just dealing with a PDB file, with no chains?
### If so, we're in the pre-find-interface-chains stage.
### Figure out which early step needs to be done and submit a job to do it.
######

# Were no chains passed?
if (!defined $chain2) {

	# If so, it must be before find_interface chains was run. 
	print STDERR "Must be before find_interface chains was run.\n";

	# Has the PDB file not been downloaded yet? 
	if (&is_pdb_downloaded($pdb) == 0) {
	
		print STDERR "Creating job to download PDB file $pdb\n";

		# Submit a job to download it.
		my @cmds_curl = ["module load perl", "perl scripts/curl_for_pdb_files.pl $pdb"];

		my $job_curl = PBS::Client::Job->new(
			name => "curl_$pdb",
			wallt => "0:10:00", 
			cmd => \@cmds_curl,
			maillist => @emails,
			mailopt => "bea",
		);
		
		my $retval = $pbs->qsub($job_curl);

	} else {

		# Otherwise, run a job to call find_interface_chains.pl. 
		
		print STDERR "Creating job to call find_interface_chains.pl on $pdb\n";

		# Submit a job to download it.
		my @cmds_findint = ["module load perl", "perl scripts/find_interface_chains.pl $pdb"];

		my $job_findint = PBS::Client::Job->new(
			name => "findint_$pdb",
			wallt => "0:10:00", 
			cmd => \@cmds_findint,
			maillist => @emails,
			mailopt => "bea",
		);
		
		my $retval = $pbs->qsub($job_findint);
	}
	
	# ...and then die.
	exit 0;
}

# If we reach here, chains were passed. 
print STDERR "Must be after find_interface chains was run.\n";

######
### Is the Rosetta status anything other than done in the database?
### If so, we're in either the pre-Rosetta or Rosetta stage.
######

# Is Rosetta status listed as anything other than "done"?
if (! &is_rosetta_done($pdb, $chain1, $chain2) ) {

	# OK, we're in either the pre-Rosetta or Rosetta stage
	print STDERR "Must be in either the pre-Rosetta or Rosetta stage.\n";

	# We check if Rosetta is running, by looking for the job ID. 
	if ( &is_rosetta_running($pdb, $chain1, $chain2) ) {

		# If it is, we change the Rosetta status and quit. 
	
		my $rosetta_sql = "UPDATE pdb_partner ";
		$rosetta_sql .= "SET rosetta_status = 'RUNNING' ";
		$rosetta_sql .= "WHERE EXISTS (";
			$rosetta_sql .= "SELECT * FROM pdb, pdb_partner ";
			$rosetta_sql .= "WHERE pdb.pdb_id = '" . $pdb . "' AND ";
			$rosetta_sql .= "pdb_partner.chain_1 = '" . $chain1 . "' AND ";
			$rosetta_sql .= "pdb_partner.chain_2 = '" . $chain2 . "' AND ";
			$rosetta_sql .= "pdb.id = pdb_partner.pdb_id);";
		
		print STDERR "Changing Rosetta status to RUNNING\n";
		print STDERR $rosetta_sql . "\n";
		
		my $dbh = DBI->connect(
			"dbi:SQLite:dbname=workflow_tracker.db",
			"",
			"",
			{ RaiseError => 1 },
		) or die $DBI::errstr;

		my $sth = $dbh->prepare($rosetta_sql);
		$sth->execute();
		$sth->finish();
	
		$dbh->disconnect();
	
		exit 0;
	
	# Otherwise, we check if Rosetta is ACTUALLY finished despite what the DB says. 
	} elsif ( &is_rosetta_finished($pdb, $chain1, $chain2) ) {
		# This will only be true if has finished successfully, 
		# but doesn't say so yet in the DB 

		print STDERR "Rosetta must be actually finished\n";

		# Change the Rosetta status and quit
		my $rosetta_sql = "UPDATE pdb_partner ";
		$rosetta_sql .= "SET rosetta_status = 'DONE' ";
		$rosetta_sql .= "WHERE EXISTS (";
			$rosetta_sql .= "SELECT * FROM pdb, pdb_partner ";
			$rosetta_sql .= "WHERE pdb.pdb_id = '" . $pdb . "' AND ";
			$rosetta_sql .= "pdb_partner.chain_1 = '" . $chain1 . "' AND ";
			$rosetta_sql .= "pdb_partner.chain_2 = '" . $chain2 . "' AND ";
			$rosetta_sql .= "pdb.id = pdb_partner.pdb_id);";
		
		print STDERR "Changing Rosetta status to DONE\n";
		print STDERR $rosetta_sql . "\n";
		
		my $dbh = DBI->connect(
			"dbi:SQLite:dbname=workflow_tracker.db",
			"",
			"",
			{ RaiseError => 1 },
		) or die $DBI::errstr;

		my $sth = $dbh->prepare($rosetta_sql);
		$sth->execute();
		$sth->finish();
	
		$dbh->disconnect();
		
		exit 0;
	
	###	# If it is, but it seems to have ended abnormally, 
	###	### Is there a test for this written yet? shouldn't it just be... Not running, but not passing &is_rosetta_finished?
	###	}
	###		# we increase the rosetta_attempts count, 
	###
	###		# change the Rosetta status, and quit. 

	# Otherwise, we check if Rosetta has failed many times, like more than three. 
	
	} elsif ( &is_rosetta_always_failing($pdb, $chain1, $chain2) ) {
	
		# If so, we change the status to "problem" and quit.
		
		print STDERR "Rosetta is always failing\n";
	
		my $rosetta_sql = "UPDATE pdb_partner ";
		$rosetta_sql .= "SET rosetta_status = 'PROBLEM' ";
		$rosetta_sql .= "WHERE EXISTS (";
			$rosetta_sql .= "SELECT * FROM pdb, pdb_partner ";
			$rosetta_sql .= "WHERE pdb.pdb_id = '" . $pdb . "' AND ";
			$rosetta_sql .= "pdb_partner.chain_1 = '" . $chain1 . "' AND ";
			$rosetta_sql .= "pdb_partner.chain_2 = '" . $chain2 . "' AND ";
			$rosetta_sql .= "pdb.id = pdb_partner.pdb_id);";
		
		print STDERR $rosetta_sql . "\n";
		
		my $dbh = DBI->connect(
			"dbi:SQLite:dbname=workflow_tracker.db",
			"",
			"",
			{ RaiseError => 1 },
		) or die $DBI::errstr;

		my $sth = $dbh->prepare($rosetta_sql);
		$sth->execute();
		$sth->finish();
	
		$dbh->disconnect();
		
		exit 0;

	# If all the above are all false...
	} else {
	
		# We must be in the pre-Rosetta stage.  
		print STDERR "Must be in pre-Rosetta stage\n";

		# Just in case it was deleted, we ask, "Was the PDB file (not) downloaded?" 
		if (&is_pdb_downloaded($pdb) == 0) {

			# If not, submit a job to download it, and then die.
			
			print STDERR "Creating job to download PDB file $pdb\n";

			# Submit a job to download it.
			my @cmds_curl = ["module load perl", "perl scripts/curl_for_pdb_files.pl $pdb"];

			my $job_curl = PBS::Client::Job->new(
				name => "curl_$pdb",
				wallt => "0:10:00", 
				cmd => \@cmds_curl,
				maillist => @emails,
				mailopt => "bea",
			);
		
			my $retval = $pbs->qsub($job_curl);
			
			exit 0;
						
		# We ask, "Is there (not) a PDB partner file?" 
		} elsif ( ! &is_partner_file_generated($pdb, $chain1, $chain2) ) { 

			# If not, we call a job to generate one and quit. 
			print STDERR "Creating job to call pdb_partner_single.pl on $pdb, ";
			print STDERR "chains $chain1 and $chain2.\n";

			my @cmds_pdb_part = ["module load perl", "perl scripts/pdb_partner_single.pl $pdb $chain1 $chain2"];

			my $job_pdb_part = PBS::Client::Job->new(
				name => "pdb_part_$pdb",
				wallt => "2:00:00", 
				cmd => \@cmds_pdb_part,
				maillist => @emails,
				mailopt => "bea",
			);
		
			my $retval = $pbs->qsub($job_pdb_part);

			exit 0;
		
		# If all these are OK... 
		} else {
			# We must be done with all the pre-Rosetta steps and ready to run Rosetta. 
			print STDERR "Must be ready to run Rosetta.\n";

			# We submit the job to do so, which is handled by a different script...
			system ("perl scripts/run_rosetta.pl $pdb $chain1 $chain2");

			# ...change the Rosetta status to "running" and...

			my $rosetta_sql = "UPDATE pdb_partner ";
				$rosetta_sql .= "SET rosetta_status = 'RUNNING' ";
				$rosetta_sql .= "WHERE EXISTS (";
			$rosetta_sql .= "SELECT * FROM pdb, pdb_partner ";
			$rosetta_sql .= "WHERE pdb.pdb_id = '" . $pdb . "' AND ";
			$rosetta_sql .= "pdb_partner.chain_1 = '" . $chain1 . "' AND ";
			$rosetta_sql .= "pdb_partner.chain_2 = '" . $chain2 . "' AND ";
			$rosetta_sql .= "pdb.id = pdb_partner.pdb_id);";
		
			print STDERR $rosetta_sql . "\n";
		
			my $dbh = DBI->connect(
				"dbi:SQLite:dbname=workflow_tracker.db",
				"",
				"",
				{ RaiseError => 1 },
			) or die $DBI::errstr;

			my $sth = $dbh->prepare($rosetta_sql);
			$sth->execute();
			$sth->finish();
			
			# ...increase the number of Rosetta attempts. 
			
			$rosetta_sql = "UPDATE pdb_partner ";
				$rosetta_sql .= "SET rosetta_attempts = rosetta_attempts + 1 ";
				$rosetta_sql .= "WHERE EXISTS (";
			$rosetta_sql .= "SELECT * FROM pdb, pdb_partner ";
			$rosetta_sql .= "WHERE pdb.pdb_id = '" . $pdb . "' AND ";
			$rosetta_sql .= "pdb_partner.chain_1 = '" . $chain1 . "' AND ";
			$rosetta_sql .= "pdb_partner.chain_2 = '" . $chain2 . "' AND ";
			$rosetta_sql .= "pdb.id = pdb_partner.pdb_id);";
		
			print STDERR $rosetta_sql . "\n";

			$sth = $dbh->prepare($rosetta_sql);
			$sth->execute();
			$sth->finish();
			
			$dbh->disconnect();
			
			exit 0;
		
		}
	
	}

	exit 0;
}

######
### If we're here, Rosetta status is "done" in the database
### Therefore, we're in the proc_ala stage.
######

print STDERR "In proc_ala stage.\n";

# We ask "Is proc_ala running?" 

# If so, we change the DB status and quit. 

# We ask, "Is proc_ala done?" 

# If so, we determine if it was successful or faulty. 

# If problematic we change db status, increment the proc_ala attempts, and quit. 

# If it was successful, we list it as "done," delete extraneous files, and clean up the DB. 

# Done! 

# Otherwise we check "has proc_ala failed many times?" 

# If so, we change the status to "problem" or something, and quit. 

# Otherwise, we must be ready to run proc_ala. 

# We submit the job to do so, changing the proc_ala status to "running" and increasing the number of proc_ala attempts. 

